# -*- coding: utf-8 -*-

"""
Comment:

"""
import numpy as np
import scipy.io
from scipy.stats import uniform
from cdlp55 import calc_upsellpossibilities
import copy

def offerset_based(segNP,mode,inst,revenues_UHbinom_run,bound,eps,checkcapconst,overview,capcon,runn,nRes,rem_cap,nProducts,offerSets,prices,eng,segQ,segScale,cap,revenues_UP_run,time_elapsed_cont,uphierar,nyp,segLambdas):
    #generate new customer segments
    uniqS, counts = np.unique(overview[runn,:,1:4],return_counts=True,axis=0)
    uniqSeg = uniqS[uniqS[:,1]!=-10].astype(int)                                         #delete nopurch rows
    counts2 = counts[uniqS[:,1]!=-10] 
    counts2 = counts2[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-1]==1),invert=True)]
    uniqSeg = uniqSeg[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-1]==1),invert=True)]   #delete those in highest compartment
    counts2 = counts2[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-1]==2),invert=True)]
    uniqSeg = uniqSeg[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-1]==2),invert=True)]   #delete those in highest compartment
    if  rem_cap[0][1:][-1] == 0:
        counts2 = counts2[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-2]==1),invert=True)]
        uniqSeg = uniqSeg[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-2]==1),invert=True)]   #delete those in highest compartment
        counts2 = counts2[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-2]==2),invert=True)]
        uniqSeg = uniqSeg[np.isin(uniqSeg[:,1],np.where(capcon[:,nRes-2]==2),invert=True)]   #delete those in highest compartment
    
    [uplist,klist,p_comp,uplist_M,y_M] = calc_upsellpossibilities(uphierar,uniqSeg,counts2,nProducts,nRes,offerSets,prices,capcon,rem_cap)
        
    #save arrays to load into matlab     
    scipy.io.savemat('data6c.mat',{'inst':inst,'segNP':segNP,'nyp':nyp,'bound':bound,'eps':eps,'Y':y_M,'upMat':uplist_M,\
                                   'os':offerSets,'uniqSeg':uniqSeg,'prices':prices,'segQ':segQ,'capcon':capcon,\
                                   'segScale':segScale,'cap':cap,'rem_cap':rem_cap,'cap':cap})
    
    #Upsell problem with all different segments = optimal price for upsell for all segments (optimal price point)
    if mode ==  '':
        [v,rev,tsub,infeas,binom,probUP,nn,timertime] = eng.UP6c_p(nargout=8)
    elif mode == '_p2':
        [v,rev,tsub,infeas,binom,probUP,nn] = eng.UP6c_p2(nargout=7)
        timertime = 0
    time_elapsed_cont[runn] = tsub
    checkcapconst[runn,0] = infeas
    revenues_UHbinom_run[runn] = binom
            
    #save values from optimizations  
    values = np.array(list(v))                  #Attention: Matlab starts with 1, Python with 0
    values[:,0:2] = values[:,0:2]-1             #Convert it back to start from 0
    revenues_UP_run[runn] = -rev
    if type(probUP) == float:
        prob = np.array([probUP])
        nns = np.array(([nn]))
    else:
        prob = np.squeeze(np.array(list(probUP)))
        nns = np.squeeze(np.array(list(nn)))   
                          
    return prob,nns,values,uniqSeg,revenues_UP_run,revenues_UHbinom_run,time_elapsed_cont,checkcapconst,timertime